use futures_util::{SinkExt, StreamExt};
use serde_json::json;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::time::sleep;
use tokio_tungstenite::{connect_async, tungstenite::Message};
use uuid::Uuid;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let url = "ws://127.0.0.1:8080";
    let num_clients = 100;
    let messages_per_client = 1000;
    let message_interval = Duration::from_millis(10);
    
    println!("WebSocket Benchmark Test");
    println!("URL: {}", url);
    println!("Number of clients: {}", num_clients);
    println!("Messages per client: {}", messages_per_client);
    println!("Message interval: {:?}", message_interval);
    println!();
    
    let total_messages = Arc::new(AtomicU64::new(0));
    let successful_messages = Arc::new(AtomicU64::new(0));
    let failed_connections = Arc::new(AtomicU64::new(0));
    
    let start_time = Instant::now();
    let mut handles = vec![];
    
    for i in 0..num_clients {
        let total_messages = Arc::clone(&total_messages);
        let successful_messages = Arc::clone(&successful_messages);
        let failed_connections = Arc::clone(&failed_connections);
        
        let handle = tokio::spawn(async move {
            match run_client(url, i, messages_per_client, message_interval).await {
                Ok(success_count) => {
                    successful_messages.fetch_add(success_count, Ordering::Relaxed);
                    total_messages.fetch_add(messages_per_client, Ordering::Relaxed);
                }
                Err(e) => {
                    eprintln!("Client {} failed: {}", i, e);
                    failed_connections.fetch_add(1, Ordering::Relaxed);
                }
            }
        });
        
        handles.push(handle);
        
        // Small delay between client connections to avoid overwhelming the server
        sleep(Duration::from_millis(10)).await;
    }
    
    // Wait for all clients to complete
    for handle in handles {
        let _ = handle.await;
    }
    
    let duration = start_time.elapsed();
    let total = total_messages.load(Ordering::Relaxed);
    let successful = successful_messages.load(Ordering::Relaxed);
    let failed = failed_connections.load(Ordering::Relaxed);
    
    println!("\nBenchmark Results:");
    println!("==================");
    println!("Total duration: {:?}", duration);
    println!("Total clients: {}", num_clients);
    println!("Successful connections: {}", num_clients - failed as usize);
    println!("Failed connections: {}", failed);
    println!("Total messages attempted: {}", total);
    println!("Successful messages: {}", successful);
    println!("Failed messages: {}", total - successful);
    println!("Success rate: {:.2}%", (successful as f64 / total as f64) * 100.0);
    println!("Messages per second: {:.2}", successful as f64 / duration.as_secs_f64());
    println!("Average latency per message: {:.2}ms", duration.as_millis() as f64 / successful as f64);
    
    Ok(())
}

async fn run_client(
    url: &str,
    client_id: usize,
    message_count: u64,
    interval: Duration,
) -> Result<u64, Box<dyn std::error::Error>> {
    let (ws_stream, _) = connect_async(url).await?;
    let (mut write, mut read) = ws_stream.split();
    
    let user_id = Uuid::new_v4();
    let username = format!("bench_user_{}", client_id);
    
    // Send join message
    let join_msg = json!({
        "type": "join",
        "user_id": user_id,
        "username": username
    });
    write.send(Message::Text(join_msg.to_string())).await?;
    
    // Spawn task to handle incoming messages
    let read_task = tokio::spawn(async move {
        let mut count = 0u64;
        while let Some(msg) = read.next().await {
            if let Ok(Message::Text(_)) = msg {
                count += 1;
            }
        }
        count
    });
    
    let mut successful_sends = 0u64;
    
    // Send messages
    for i in 0..message_count {
        let message = json!({
            "type": "broadcast",
            "content": format!("Benchmark message {} from client {}", i, client_id),
            "sender_id": user_id,
            "timestamp": current_timestamp()
        });
        
        match write.send(Message::Text(message.to_string())).await {
            Ok(_) => {
                successful_sends += 1;
            }
            Err(e) => {
                eprintln!("Client {} failed to send message {}: {}", client_id, i, e);
                break;
            }
        }
        
        sleep(interval).await;
    }
    
    // Wait a bit for responses
    sleep(Duration::from_secs(2)).await;
    
    // Close connection
    let _ = write.close().await;
    
    // Get received message count
    let received_count = read_task.await.unwrap_or(0);
    
    Ok(std::cmp::min(successful_sends, received_count))
}

fn current_timestamp() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs()
}